package main

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"os"
)

func main() {
	// 生成 RSA 密钥对
	privateKey, publicKey := generateRSAKeyPair(2048)

	// 要加密的数据
	data := []byte("123456")

	// 使用公钥加密数据
	encryptedData, err := publicKeyEncrypt(publicKey, data)
	print("加密后的密文： ", string(encryptedData))
	if err != nil {
		fmt.Println("Encryption failed:", err)
		return
	}
	fmt.Println("Encrypted Data:", encryptedData)

	// 使用私钥解密数据
	decryptedData, err := privateKeyDecrypt(privateKey, encryptedData)
	if err != nil {
		fmt.Println("Decryption failed:", err)
		return
	}
	fmt.Println("Decrypted Data:", string(decryptedData))
}

// 生成 RSA 密钥对
func generateRSAKeyPair(bits int) (*rsa.PrivateKey, *rsa.PublicKey) {
	privateKey, err := rsa.GenerateKey(rand.Reader, bits)
	if err != nil {
		panic(err)
	}

	// 提取公钥
	publicKey := &privateKey.PublicKey

	// 可选：保存密钥到文件
	savePrivateKey("private_key.pem", privateKey)
	savePublicKey("public_key.pem", publicKey)

	return privateKey, publicKey
}

// 使用公钥加密数据
func publicKeyEncrypt(publicKey *rsa.PublicKey, data []byte) ([]byte, error) {
	return rsa.EncryptPKCS1v15(rand.Reader, publicKey, data)
}

// 使用私钥解密数据
func privateKeyDecrypt(privateKey *rsa.PrivateKey, encryptedData []byte) ([]byte, error) {
	return rsa.DecryptPKCS1v15(rand.Reader, privateKey, encryptedData)
}

// 保存私钥到文件
func savePrivateKey(filename string, privateKey *rsa.PrivateKey) {
	privateKeyBytes := x509.MarshalPKCS1PrivateKey(privateKey)
	privateKeyPEM := pem.EncodeToMemory(
		&pem.Block{
			Type:  "RSA PRIVATE KEY",
			Bytes: privateKeyBytes,
		},
	)

	err := os.WriteFile(filename, privateKeyPEM, 0600)
	if err != nil {
		panic(err)
	}
	fmt.Println("Private key saved to", filename)
}

// 保存公钥到文件
func savePublicKey(filename string, publicKey *rsa.PublicKey) {
	publicKeyBytes, err := x509.MarshalPKIXPublicKey(publicKey)
	if err != nil {
		panic(err)
	}

	publicKeyPEM := pem.EncodeToMemory(
		&pem.Block{
			Type:  "PUBLIC KEY",
			Bytes: publicKeyBytes,
		},
	)

	err = os.WriteFile(filename, publicKeyPEM, 0644)
	if err != nil {
		panic(err)
	}
	fmt.Println("Public key saved to", filename)
}
